import torch
import torch.nn as nn

from activation import *
from pos_encoding import *

#添加边缘卷积类
class EdgeConv(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(EdgeConv, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_channels * 2, out_channels, kernel_size=1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.LeakyReLU(negative_slope=0.2)
        )

    def forward(self, x, k=20):
        # 构建k近邻图
        pairwise_distance = -2 * torch.matmul(x, x.transpose(2, 1)) + torch.sum(x ** 2, dim=2,
                                                                                keepdim=True) - torch.sum(x ** 2, dim=2,
                                                                                                          keepdim=True).transpose(
            2, 1)
        idx = pairwise_distance.topk(k=k, dim=-1)[1]

        # 获取k近邻特征并计算边特征
        idx_base = torch.arange(0, x.size(0), device=x.device).view(-1, 1, 1) * x.size(1)
        idx = idx + idx_base
        feature = x.view(x.size(0) * x.size(1), -1)[idx, :].view(x.size(0), x.size(1), k, x.size(2))
        edge_features = torch.cat((feature - x.unsqueeze(2), x.unsqueeze(2)), dim=3)

        # 应用卷积并聚合特征
        edge_features = edge_features.permute(0, 3, 1, 2)
        out = self.conv(edge_features)
        out = torch.max(out, dim=-1)[0].permute(0, 2, 1)

        return out


class DGCNN_MLP(nn.Module):
    def __init__(self, size, k=20, activation_type='relu', bias=True, num_pos_encoding=-1):
        super(DGCNN_MLP, self).__init__()
        self.size = size
        self.k = k
        self.num_layer = len(size) - 1
        self.bias = bias

        self.layers = nn.ModuleList()

        # 1. 添加位置编码（保留原始功能）
        if num_pos_encoding > 0:
            pos_encoder = PosEncoding(d_in=size[0], num_freqs=num_pos_encoding)
            self.layers.append(pos_encoder)
            size[0] = pos_encoder.d_out

        # 2. 使用EdgeConv层替代线性层
        for i in range(self.num_layer - 1):  # 最后一层不使用EdgeConv
            self.layers.append(EdgeConv(size[i], size[i + 1], k=self.k))

            # 添加激活函数（保持与原MLP一致）
            if activation_type == 'relu':
                self.layers.append(nn.ReLU(inplace=True))
            elif activation_type == 'lrelu':
                self.layers.append(nn.LeakyReLU(0.2, inplace=True))
            elif activation_type == 'sine':
                self.layers.append(Sine())

        # 3. 最后一层使用线性层（用于输出最终维度）
        self.layers.append(nn.Linear(size[-2], size[-1], bias=bias))

    def forward(self, x):
        # x: (batch_size, num_points, in_channels)
        for layer in self.layers:
            x = layer(x)
        return x

    def forward_simple(self, x):
        return self.net(x)

    def layer_feature(self, x, k):
        return self.net[:k](x)


class MultiScale_DGCNN_MLP(nn.Module):
    def __init__(self, size, k=20, activation_type='relu', num_pos_encoding=-1):
        super().__init__()
        self.layers = nn.ModuleList()
        self.feature_dims = []

        # 添加位置编码
        if num_pos_encoding > 0:
            self.layers.append(PosEncoding(d_in=size[0], num_freqs=num_pos_encoding))
            size[0] = self.layers[0].d_out

        # 添加多个EdgeConv层并保存中间特征
        for i in range(len(size) - 2):
            self.layers.append(EdgeConv(size[i], size[i + 1]))
            self.feature_dims.append(size[i + 1])
            self.layers.append(getattr(nn, activation_type)(inplace=True))

        # 最后一层融合所有尺度的特征
        self.layers.append(nn.Linear(sum(self.feature_dims), size[-1]))

    def forward(self, x):
        features = []
        for layer in self.layers:
            if isinstance(layer, EdgeConv):
                x = layer(x)
                features.append(x)  # 保存中间特征
            elif isinstance(layer, nn.Linear):
                x = torch.cat(features, dim=-1)  # 融合多尺度特征
                x = layer(x)
            else:
                x = layer(x)
        return x



if __name__ == '__main__':
    mlp = MLP((3, 4, 5))
    print(mlp)
    print('num of param: ', len(list(mlp.parameters())))
    a = torch.rand(10, 3)
    b = mlp(a)
    print(b.shape)